import os
from PIL import Image, ImageDraw, ImageFont

class_id_to_enname = {
    0: "part_damaged_ship",
    1: "front_damaged_launchcar",
    2: "back_damaged_launchcar",
    3: "hard_damaged_launchcar",
    4: "turnover_launchcar",
    5: "part_damaged_plane",
    6: "hard_damaged_plane",
    7: "fuel_leaking_plane",
    8: "part_damaged_radarcar",
    9: "hard_damaged_radarcar",
    10: "turnover_radarcar",
    11: "ship",
    12: "launch_car",
    13: "radar_car",
    14: "plane"
}

def parse_yolo_labels(label_path, img_width, img_height):
    """解析YOLO格式的标签并将相对坐标转换为像素坐标"""
    boxes = []
    with open(label_path, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) == 5:
                cls, x, y, w, h = map(float, parts)
                # 将相对坐标转换为像素坐标
                x_min = (x - w / 2) * img_width
                y_min = (y - h / 2) * img_height
                x_max = (x + w / 2) * img_width
                y_max = (y + h / 2) * img_height
                boxes.append((cls, x_min, y_min, x_max, y_max))
    return boxes


def get_color_from_class(cls):
    # YOLO标准16色色板（与网页2、网页4配色一致）
    yolo_color_palette = [
        (0, 0, 0),       # 0. 黑色
        (0, 0, 255),     # 1. 红色
        (0, 255, 0),     # 2. 绿色
        (255, 0, 0),     # 3. 蓝色
        (0, 255, 255),   # 4. 黄色
        (255, 165, 0),   # 5. 橙色
        (75, 0, 130),    # 6. 靛蓝
        (255, 20, 147),  # 7. 深粉色
        (0, 255, 127),   # 8. 春绿色
        (255, 255, 0),   # 9. 青色
        (128, 0, 128),   # 10. 紫色
        (255, 99, 71),   # 11. 番茄红
        (60, 179, 113),  # 12. 海洋绿
        (238, 130, 238), # 13. 紫罗兰
        (255, 215, 0),   # 14. 金色
        (0, 191, 255)    # 15. 天空蓝
    ]
    return yolo_color_palette[cls % len(yolo_color_palette)]


def draw_boxes(image_path, label_path, output_folder):
    """根据标签绘制边界框到图像上，并保存图像"""

    img = Image.open(image_path)
    if img.mode in ("P", "L", "1"):
        img = img.convert("RGB")
    img_width, img_height = img.size

    boxes = parse_yolo_labels(label_path, img_width, img_height)

    draw = ImageDraw.Draw(img)
    # 加载字体文件（如arial.ttf），设置字号为30（默认字号通常为10-12）
    font = ImageFont.truetype("arial.ttf", size=30)
    # 绘制每个框
    for box in boxes:
        cls, x_min, y_min, x_max, y_max = box
        cls = int(cls)  # 转换为整数标签
        color = get_color_from_class(cls)  # 为每个类别生成颜色
        en_name = class_id_to_enname.get(cls, f"unknown id")
        draw.rectangle([x_min, y_min, x_max, y_max], outline=color, width=2)
        draw.text((x_min, y_min), f"{en_name}", fill=color, font=font)

    image_name = os.path.basename(image_path)
    save_path = os.path.join(output_folder, image_name)

    img.save(save_path)
    print(f"✅ 图像已保存: {save_path}")


def draw_image_bylabel(image_folder, label_folder, output_folder):
    """
    根据给出的img和 label文件夹将yolo的框绘制到图中
    :param image_folder: 图片文件夹
    :param label_folder: 标签文件夹
    :param output_folder: 标注后的文件夹
    :return:
    """
    os.makedirs(output_folder, exist_ok=True)
    extensions = ('jpg', 'jpeg', 'png', 'webp', 'bmp')
    image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(extensions)]

    for image_file in image_files:
        image_path = os.path.join(image_folder, image_file)
        label_path = os.path.join(label_folder, os.path.splitext(image_file)[0] + '.txt')

        if os.path.exists(label_path):
            draw_boxes(image_path, label_path, output_folder)


if __name__ == "__main__":
    # 函数一
    image_folder = r"F:\yolo_summary\yolo_code\car\521\train\train"  # 图片所在文件夹
    label_folder = r"F:\yolo_summary\yolo_code\car\521\train\train_label"  # 标签文件夹
    output_folder = r"F:\yolo_summary\yolo_code\car\521\train\show_result"  # 保存结果的文件夹
    draw_image_bylabel(image_folder, label_folder, output_folder)

